"""Feature extraction from Zeek SSL logs."""

import pandas as pd
import numpy as np
from typing import Dict, List, Optional
import re


class TLSFeatureExtractor:
    """Extract features from Zeek SSL/TLS logs for machine learning."""
    
    def __init__(self):
        self.cipher_strength = {
            'TLS_AES_128_GCM_SHA256': 3,
            'TLS_AES_256_GCM_SHA384': 4,
            'TLS_CHACHA20_POLY1305_SHA256': 4,
            'TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256': 3,
            'TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384': 4,
            'TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA256': 2,
            'TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA384': 3,
            'TLS_RSA_WITH_AES_128_GCM_SHA256': 2,
            'TLS_RSA_WITH_AES_256_GCM_SHA384': 3,
            'TLS_RSA_WITH_AES_128_CBC_SHA256': 1,
            'TLS_RSA_WITH_AES_256_CBC_SHA256': 2,
        }
        
        self.version_strength = {
            'SSLv3': 1,
            'TLSv10': 1,
            'TLSv11': 2,
            'TLSv12': 3,
            'TLSv13': 4,
        }
    
    def parse_ssl_log(self, log_path: str) -> pd.DataFrame:
        """Parse Zeek SSL log file."""
        data = []
        
        with open(log_path, 'r') as f:
            for line in f:
                line = line.strip()
                if line.startswith('#') or not line:
                    continue
                
                fields = line.split('\t')
                if len(fields) < 17:
                    continue
                
                try:
                    record = {
                        'ts': float(fields[0]) if fields[0] != '-' else 0,
                        'uid': fields[1],
                        'id.orig_h': fields[2],
                        'id.orig_p': int(fields[3]) if fields[3] != '-' else 0,
                        'id.resp_h': fields[4],
                        'id.resp_p': int(fields[5]) if fields[5] != '-' else 0,
                        'version': fields[6],
                        'cipher': fields[7],
                        'curve': fields[8],
                        'server_name': fields[9],
                        'resumed': fields[10] == 'T',
                        'last_alert': fields[11],
                        'next_protocol': fields[12],
                        'established': fields[13] == 'T',
                        'ssl_history': fields[14],
                        'cert_chain_fps': fields[15],
                        'client_cert_chain_fps': fields[16],
                        'sni_matches_cert': fields[17] if len(fields) > 17 else 'F',
                    }
                    data.append(record)
                except (ValueError, IndexError) as e:
                    print(f"Skipping malformed line: {e}")
                    continue
        
        return pd.DataFrame(data)
    
    def extract_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """Extract ML features from parsed SSL log data."""
        features = df.copy()
        
        # Handle different column naming conventions
        # Some pandas versions might convert dots to underscores or other formats
        port_orig_col = None
        port_resp_col = None
        
        for col in features.columns:
            if 'orig_p' in col and 'port' not in col:
                port_orig_col = col
            elif 'resp_p' in col and 'port' not in col:
                port_resp_col = col
        
        if port_orig_col is None:
            # Try to find the original port column
            possible_orig_cols = ['id.orig_p', 'id_orig_p', 'orig_p']
            for col in possible_orig_cols:
                if col in features.columns:
                    port_orig_col = col
                    break
        
        if port_resp_col is None:
            # Try to find the response port column
            possible_resp_cols = ['id.resp_p', 'id_resp_p', 'resp_p']
            for col in possible_resp_cols:
                if col in features.columns:
                    port_resp_col = col
                    break
        
        if port_orig_col is None or port_resp_col is None:
            raise ValueError(f"Could not find port columns. Available columns: {list(features.columns)}")
        
        # Port-based features
        features['src_port_category'] = pd.cut(
            features[port_orig_col], 
            bins=[0, 1023, 49151, 65535], 
            labels=['well_known', 'registered', 'dynamic']
        )
        features['dst_port_is_443'] = (features[port_resp_col] == 443).astype(int)
        features['dst_port_is_8443'] = (features[port_resp_col] == 8443).astype(int)
        
        # TLS version features
        features['tls_version_strength'] = features['version'].map(self.version_strength).fillna(0)
        features['is_tls13'] = (features['version'] == 'TLSv13').astype(int)
        features['is_sslv3'] = (features['version'] == 'SSLv3').astype(int)
        
        # Cipher suite features
        features['cipher_strength'] = features['cipher'].map(self.cipher_strength).fillna(0)
        features['is_gcm_cipher'] = features['cipher'].str.contains('GCM').astype(int)
        features['is_cbc_cipher'] = features['cipher'].str.contains('CBC').astype(int)
        features['is_chacha_cipher'] = features['cipher'].str.contains('CHACHA').astype(int)
        
        # Server name features
        features['has_server_name'] = (features['server_name'] != '-').astype(int)
        features['server_name_length'] = features['server_name'].apply(
            lambda x: len(x) if x != '-' else 0
        )
        features['server_name_has_subdomain'] = features['server_name'].apply(
            lambda x: x.count('.') > 1 if x != '-' else 0
        )
        features['server_name_is_ip'] = features['server_name'].apply(
            lambda x: bool(re.match(r'^\d+\.\d+\.\d+\.\d+$', x)) if x != '-' else 0
        )
        
        # Certificate features
        features['cert_chain_count'] = features['cert_chain_fps'].apply(
            lambda x: len(x.split(',')) if x != '(empty)' and x != '-' else 0
        )
        features['has_client_cert'] = (features['client_cert_chain_fps'] != '(empty)').astype(int)
        
        # Connection features
        features['connection_established'] = features['established'].astype(int)
        features['session_resumed'] = features['resumed'].astype(int)
        
        # SSL history features
        features['ssl_history_length'] = features['ssl_history'].apply(len)
        features['has_ssl_history'] = (features['ssl_history'] != '-').astype(int)
        
        # Alert features
        features['has_alert'] = (features['last_alert'] != '-').astype(int)
        
        # Next protocol features
        features['has_next_protocol'] = (features['next_protocol'] != '-').astype(int)
        features['is_http_protocol'] = features['next_protocol'].str.contains('http').astype(int)
        
        # SNI matching certificate
        features['sni_matches_cert_bool'] = (features['sni_matches_cert'] == 'T').astype(int)
        
        return features
    
    def get_feature_columns(self) -> List[str]:
        """Get list of feature columns for ML model."""
        return [
            'src_port_category', 'dst_port_is_443', 'dst_port_is_8443',
            'tls_version_strength', 'is_tls13', 'is_sslv3',
            'cipher_strength', 'is_gcm_cipher', 'is_cbc_cipher', 'is_chacha_cipher',
            'has_server_name', 'server_name_length', 'server_name_has_subdomain', 'server_name_is_ip',
            'cert_chain_count', 'has_client_cert',
            'connection_established', 'session_resumed',
            'ssl_history_length', 'has_ssl_history',
            'has_alert', 'has_next_protocol', 'is_http_protocol',
            'sni_matches_cert_bool'
        ]